import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from data_process import plot_clusters


def pod_reduce_and_restore(y: pd.DataFrame, target_dimension: int = 1):
    """
    Perform POD to reduce the dimension of the input data and restore it.

    Parameters:
    y (pd.DataFrame): Multi-dimensional data to be reduced (columns are dimensions).
    target_dimension (int): Dimension to reduce to. Default is 1.

    Returns:
    tuple: eigenvalues (reduced dimension result), eigenvectors (used for restoration), and mean vector.
    """
    # Step 1: Convert DataFrame to numpy array and mean-center
    y_array = y.values
    mean_vector = np.mean(y_array, axis=0)
    mean_centered_data = y_array - mean_vector

    # Step 2: Perform Singular Value Decomposition (SVD)
    U, S, VT = np.linalg.svd(mean_centered_data, full_matrices=False)

    # Step 3: Extract the reduced dimensions and eigenvectors
    reduced_dimension = S[:target_dimension] * U[:, :target_dimension]
    eigenvectors = VT[:target_dimension, :]

    return reduced_dimension, eigenvectors, mean_vector


def calculate_eigenvalue_and_restore(test_data: pd.DataFrame, eigenvectors: np.ndarray, mean_vector: np.ndarray):
    """
    Calculate eigenvalues for test data and restore it using eigenvectors and mean vector.

    Parameters:
    test_data (pd.DataFrame): Test data to calculate eigenvalues and restore.
    eigenvectors (np.ndarray): Eigenvectors obtained from training data.
    mean_vector (np.ndarray): Mean vector obtained from training data.

    Returns:
    tuple: eigenvalues (reduced dimension result) and restored data.
    """
    test_array = test_data.values
    mean_centered_test_data = test_array - mean_vector

    # Calculate eigenvalues
    eigenvalues = np.dot(mean_centered_test_data, eigenvectors.T)

    # Restore the data
    restored_data = np.dot(eigenvalues, eigenvectors) + mean_vector

    return eigenvalues, restored_data



if __name__ == "__main__":
    # Load the data from the file
    file_path = './data/sample1.xlsx'
    data = pd.read_excel(file_path, sheet_name='Sheet3')

    # Separate x (Force) and y (Node data)
    x = data['Force']
    y = data.loc[:, "Node1":"Node2900"]

    # # Split the data into two parts
    # y_train = y.iloc[:len(y)//2]
    # y_test = y.iloc[len(y)//2:]

    # Split the data into training and testing sets
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.5, random_state=42)

    # Use training data to calculate eigenvectors and mean
    eigenvalues_train, eigenvectors, mean_vector = pod_reduce_and_restore(y_train, target_dimension=3)

    plot_clusters(pd.DataFrame(x_train), pd.DataFrame(eigenvalues_train))

    # Use test data to calculate eigenvalues and restore the data
    eigenvalues_test, restored_test_data = calculate_eigenvalue_and_restore(y_test, eigenvectors, mean_vector)

    # Save the eigenvalues, eigenvectors, and mean_vector
    eigenvalues_train_df = pd.DataFrame(eigenvalues_train, columns=[f'Component_{i+1}' for i in range(eigenvalues_train.shape[1])])
    result_df = pd.concat([pd.DataFrame(x_train.values, columns=["Force"]), eigenvalues_train_df], axis=1)
    result_df.to_csv('./output/reduced_data_train.csv', index=False)

    eigenvalues_test_df = pd.DataFrame(eigenvalues_test, columns=[f'Component_{i+1}' for i in range(eigenvalues_test.shape[1])])
    result_df = pd.concat([pd.DataFrame(x_test.values, columns=["Force"]), eigenvalues_test_df], axis=1)
    result_df.to_csv('./output/reduced_data_test.csv', index=False)

    eigenvector_df = pd.DataFrame(eigenvectors, columns=[f'Node_{i+1}' for i in range(eigenvectors.shape[1])])
    eigenvector_df.to_csv('./output/eigenvectors.csv', index=False)

    eigenvector_df = pd.DataFrame(mean_vector, columns=["Mean_Vector"])
    eigenvector_df.to_csv('./output/mean_vector.csv', index=False)

    # Plot the difference between the restored data and the original test data
    difference = restored_test_data - y_test.values
    plt.figure(figsize=(10, 6))
    plt.plot(np.mean(np.abs(difference), axis=0), label='Mean Absolute Difference')
    plt.xlabel('Node Index')
    plt.ylabel('Difference')
    plt.title('Difference Between Restored and Original Test Data')
    plt.legend()
    plt.grid()
    plt.savefig('./output/difference_plot_test.png')
    plt.show()

    # Plot the original and restored data for comparison
    plt.figure(figsize=(10, 6))
    plt.plot(np.mean(y_test.values, axis=0), label='Original Test Data', alpha=0.7)
    plt.plot(np.mean(restored_test_data, axis=0), label='Restored Test Data', alpha=0.7)
    plt.xlabel('Node Index')
    plt.ylabel('Value')
    plt.title('Comparison of Original and Restored Test Data')
    plt.legend()
    plt.grid()
    plt.savefig('./output/comparison_plot_test.png')
    plt.show()
